import os
import numpy as np

os.chdir(os.path.dirname(__file__))


def info2digit_world(InfoList, tokenizer, max_symbols, max_lines):
    instances = InfoList
    # instances to batch ids
    true_labels = []
    input_ids_s = []
    masks = []
    matrixs = []
    formulas = []
    rel_type2weight = {2: 0.512,
                       3: 0.141,
                       4: 0.122,
                       5: 0.112,
                       6: 0.061,
                       7: 0.043,
                       8: 0.006,
                       9: 0.003}
    assert sum(rel_type2weight.values()) == 1
    for instance in instances:
        formulas.append(instance["公式:"])
        true_label = np.zeros(max_symbols, dtype=int)
        true_label[1:1 + instance["len_split"]] = instance["标签:"]
        true_labels.append(true_label)
        input_ids = tokenizer.encode(formulas[-1], max_length=max_symbols, truncation=True, padding="max_length",
                                     add_special_tokens=True)
        input_ids_s.append(np.array(input_ids))
        mask = np.array(input_ids)
        mask[:instance["len_split"] + 2] = 1
        masks.append(np.array(mask))
        matrix = np.eye(max_symbols, dtype=float)
        for rel in instance["relationships"]:
            matrix[rel[:2]] = rel_type2weight[rel[2]]
        matrixs.append(matrix)
    all_data = list(zip(matrixs, input_ids_s, masks, true_labels, formulas))
    return all_data